from PIL import Image
import torchvision
import torch
from torch import nn


class Demo(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, 1, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64 * 4 * 4, 64),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        return self.model(x)


transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((32, 32)),
    torchvision.transforms.ToTensor()
])

ans_list = [('airplane', '飞机'), ('automobile', '汽车'), ('bird', '鸟'), ('cat', '猫'), ('deer', '鹿'), ('dog', '狗'),
            ('frog', '青蛙'), ('horse', '马'), ('ship', '船'), ('truck', '卡车')]


def deep_model_check_image_2(image_path):
    image = Image.open(image_path).convert('RGB')  # png 有四个通道

    image = transform(image)

    # 加载模型权重
    state_dict = torch.load('./demo_29.pth', map_location=torch.device('cpu'), weights_only=True)
    # 创建模型实例
    model = Demo()
    # 加载权重到模型实例
    model.load_state_dict(state_dict)

    image = image.reshape([1, 3, 32, 32])

    model.eval()
    with torch.no_grad():
        output = model(image)

    return ans_list[output.argmax(1)]


if __name__ == "__main__":
    image_path = 'test_img/dog_2.png'
    print(deep_model_check_image_2(image_path))
